# tasks/file_tasks.py

import os
import uuid
import json
import traceback
import time
from typing import List, Dict, Any
from celery import current_task

from config.settings import settings
from config.model_config import ModelConfig
from tasks.celery_app import celery_app
from app.core.database import connect_milvus, init_collections, get_text_collection, get_image_collection
from app.core.logging import get_logger
from app.processors.text_processor import TextProcessor
from app.processors.image_processor import ImageProcessor
from app.processors.video_processor import VideoProcessor
from app.processors.document_processor import DocumentProcessor

logger = get_logger(__name__)

@celery_app.task(bind=True, queue='file_processing')
def process_file_task(self, file_path: str, file_type: str, task_id: str):
    """处理单个文件的异步任务"""
    start_time = time.time()
    
    try:
        # 更新任务状态
        self.update_state(
            state=settings.TASK_STATUS_PROCESSING,
            meta={'current': 0, 'total': 100, 'status': '开始处理文件...'}
        )
        
        logger.info(f"开始处理文件: {file_path} (类型: {file_type})")
        
        # 初始化数据库
        init_collections()
        
        # 根据文件类型选择处理器
        if file_type in ['pdf', 'docx', 'pptx']:
            processor = DocumentProcessor()
        elif file_type in ['jpg', 'jpeg', 'png', 'bmp', 'gif']:
            processor = ImageProcessor()
        elif file_type in ['mp4', 'avi', 'mov', 'mkv']:
            processor = VideoProcessor()
        else:
            processor = TextProcessor()
        
        # 处理文件
        result = processor.process(file_path, file_type)
        
        # 更新任务状态
        final_result = {
            'status': settings.TASK_STATUS_COMPLETED,
            'file_path': file_path,
            'file_type': file_type,
            'text_chunks_count': result.get('text_chunks_count', 0),
            'images_count': result.get('images_count', 0),
            'processing_time': time.time() - start_time,
            'message': '文件处理完成'
        }
        
        self.update_state(
            state=settings.TASK_STATUS_COMPLETED,
            meta=final_result
        )
        
        logger.info(f"文件处理完成: {file_path} - 耗时: {final_result['processing_time']:.2f}s")
        
        return final_result
        
    except Exception as e:
        error_msg = f"处理文件 {file_path} 时发生错误: {str(e)}"
        logger.error(error_msg)
        logger.error(traceback.format_exc())
        
        result = {
            'status': settings.TASK_STATUS_FAILED,
            'file_path': file_path,
            'file_type': file_type,
            'error': str(e),
            'message': error_msg,
            'processing_time': time.time() - start_time
        }
        
        self.update_state(
            state=settings.TASK_STATUS_FAILED,
            meta=result
        )
        
        return result
    
    finally:
        # 断开数据库连接
        from pymilvus import connections
        if "default" in connections.list_connections():
            connections.disconnect("default")

@celery_app.task(bind=True, queue='batch_processing')
def process_batch_files_task(self, file_paths: List[str], task_id: str):
    """批量处理文件的异步任务"""
    start_time = time.time()
    
    try:
        total_files = len(file_paths)
        completed_files = 0
        results = []
        
        logger.info(f"开始批量处理 {total_files} 个文件")
        
        self.update_state(
            state=settings.TASK_STATUS_PROCESSING,
            meta={'current': 0, 'total': total_files, 'status': '开始批量处理...'}
        )
        
        for i, file_path in enumerate(file_paths):
            # 获取文件类型
            file_ext = os.path.splitext(file_path)[1].lower()
            file_type = file_ext[1:] if file_ext else 'unknown'
            
            # 处理单个文件
            result = process_file_task.delay(file_path, file_type, f"{task_id}_file_{i}")
            file_result = result.get()  # 等待完成
            
            results.append(file_result)
            completed_files += 1
            
            # 更新进度
            progress = int((completed_files / total_files) * 100)
            self.update_state(
                state=settings.TASK_STATUS_PROCESSING,
                meta={
                    'current': completed_files,
                    'total': total_files,
                    'progress': progress,
                    'status': f'已处理 {completed_files}/{total_files} 个文件'
                }
            )
        
        # 任务完成
        final_result = {
            'status': settings.TASK_STATUS_COMPLETED,
            'total_files': total_files,
            'completed_files': completed_files,
            'results': results,
            'processing_time': time.time() - start_time,
            'message': '批量文件处理完成'
        }
        
        self.update_state(
            state=settings.TASK_STATUS_COMPLETED,
            meta=final_result
        )
        
        logger.info(f"批量处理完成: {completed_files}/{total_files} 个文件 - 耗时: {final_result['processing_time']:.2f}s")
        
        return final_result
        
    except Exception as e:
        error_msg = f"批量处理文件时发生错误: {str(e)}"
        logger.error(error_msg)
        logger.error(traceback.format_exc())
        
        result = {
            'status': settings.TASK_STATUS_FAILED,
            'error': str(e),
            'message': error_msg,
            'processing_time': time.time() - start_time
        }
        
        self.update_state(
            state=settings.TASK_STATUS_FAILED,
            meta=result
        )
        
        return result 